import traceback

import paramiko #  3.4.0
import time
from io import BytesIO


class SSHClient:
    def __init__(self, host, username, password, port=22, max_retries=3, retry_delay=5):
        self.host = host
        self.username = username
        self.password = password
        self.port = port
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.client = None
        self.transport = None
        self.sftp = None

    def _connect_retry(self):
        """带重连尝试的连接方法"""
        for attempt in range(self.max_retries):
            try:
                self.client = paramiko.SSHClient()# 创建SSH客户端
                self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                self.client.connect(hostname=self.host, username=self.username, password=self.password)# 连接到远程主机

                self.transport = paramiko.Transport((self.host, self.port))
                self.transport.connect(username=self.username, password=self.password)
                # self.client = paramiko.SSHClient.from_transport(self.transport) # ssh (基于tcp)
                self.sftp = paramiko.SFTP.from_transport(self.transport) # sftp (基于ssh)
                return True
            except Exception as e:
                print(f"连接尝试 {attempt + 1}/{self.max_retries} 失败，{str(e)}")
                print(traceback.format_exc())
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay)
                else:
                    return False

    def connect(self):
        """建立SSH连接，带有重连机制"""
        if not self._connect_retry():
            raise Exception("连接失败，已达到最大重试次数")

    def disconnect(self):
        """断开SSH连接"""
        if self.sftp:
            self.sftp.close()
        if self.transport:
            self.transport.close()

    def execute_command(self, command):
        """执行远程命令并返回输出，自动重连"""
        try:
            stdin, stdout, stderr = self.client.exec_command(command)
            exit_status = stdout.channel.recv_exit_status()
            output = stdout.read().decode('utf-8')
            error = stderr.read().decode('utf-8')
            return output, error, exit_status
        except Exception as e:
            print(f"执行命令时发生错误: {str(e)}，尝试重连...")
            print(traceback.format_exc())
            self.disconnect()
            if not self._connect_retry(): # n次
                raise Exception("命令执行失败，连接重试失败")


    def upload_file(self, local_path, remote_path):
        """上传文件到远程，自动重连"""

        try:
            with open(local_path, 'rb') as f:
                self.sftp.putfo(f, remote_path)
            print(f"文件已上传: {local_path} -> {remote_path}")
            return
        except Exception as e:
            print(f"上传文件时发生错误: {str(e)}，尝试重连...")
            print(traceback.format_exc())
            self.disconnect()
            if not self._connect_retry():
                raise Exception("文件上传失败，连接重试失败")


    def download_file(self, remote_path, local_path):
        """下载远程文件到本地，自动重连"""
        retries = self.max_retries
        while retries > 0:
            try:
                with open(local_path, 'wb') as f:
                    self.sftp.getfo(remote_path, f)
                print(f"文件已下载: {remote_path} -> {local_path}")
                return
            except Exception as e:
                print(f"下载文件时发生错误: {str(e)}，尝试重连...")
                self.disconnect()
                if self._connect_retry():
                    retries -= 1
                    time.sleep(self.retry_delay)
                else:
                    raise Exception("重连失败")
        raise Exception("文件下载失败，达到最大重试次数")


    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.disconnect()


# 使用示例
if __name__ == "__main__":
    ssh_client = SSHClient('192.168.6.14', 'ps', '1')
    with ssh_client as client:  # 结束后执行SSHClient类 的__exit__()方法，自动关闭连接
        print(client.execute_command('cat /home/ps/zhangxiancai/data/231207huoni/trainV8Seg_screw/models/yolov8mSeg_640_screw2/args.yaml'))
        client.upload_file(r'D:\data\231207huoni\采图\cable\20240511\img_3_20240511_142148_415_yd_cable_NG_img.jpg', '/home/ps/zhangxiancai/data/temp/2.jpg')
        client.download_file('/home/ps/zhangxiancai/data/temp/1.jpg', r'D:\data\231207huoni\采图\1.jpg')
